import numpy as np
'''
三硬币问题：
假设有三枚硬币，分别记为A，B，C。这些硬币正面出现的概率分别是
pi,p,q,进行如下试验:先投掷硬币A，根据其结果选择硬币B或C：正面
选B，反面选C；然后掷出硬币，掷硬币结果，正面记为1，反面记为0，
独立进行n次实验，这里（n=10),观测结果如下：
1, 1, 0, 1, 0, 0, 1, 0, 1, 1。
EM法估算pi,p,q.
'''
result = np.array([1, 1, 0, 1, 0, 0, 1, 0, 1, 1,])
# 初始化参数
pi = 0.5
p = 0.5
q = 0.1


def pfunc(y, pi=pi, p=p, q=q):
    '''

    :param y: 观测结果，取值0，1，可以是一个numpy.array模型
    :param pi:
    :param p:
    :param q:
    :return: 结果是y的概率。如果传入的是一个ndarray对象，则返回一个ndarray对象
    '''
    return pi * pow(p, y) * pow((1 - p), (1 - y)) + (1 - pi) * pow(q, y) * pow((1 - q), (1 - y))


def bfunc(y, pi=pi, p=p, q=q):
    '''

    :param y: 观测结果，取值0，1，可以是一个numpy.array模型
    :param pi:
    :param p:
    :param q:
    :return: 结果是掷硬币B的概率。如果传入的是一个ndarray对象，则返回一个ndarray对象
    '''
    return pi * pow(p, y) * pow((1 - p), (1 - y)) / pfunc(y, pi=pi, p=p, q=q)


for _ in range(10):
    mid_pi = bfunc(result, pi, p, q)
    pi = np.average(mid_pi)
    p = np.dot(mid_pi, result) / np.sum(mid_pi)
    q = np.dot((1 - mid_pi), result) / np.sum(1 - mid_pi)
    print(pi, p, q)
